#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author: 邵奈一
@Email: shaonaiyi@163.com
@Date: 2024/11/13
@微信：shaonaiyi888
@微信公众号: 邵奈一 
"""
# 代码3-35
# 定义画图工具
import torch  # 导入PyTorch库
from torch import nn  # 导入PyTorch的神经网络模块
from torch.autograd import Variable  # 导入PyTorch的自动求导模块
import torchvision.transforms as tfs  # 导入PyTorch的图像处理模块
from torch.utils.data import DataLoader, sampler  # 导入PyTorch的数据加载和采样模块
from torchvision.datasets import MNIST  # 导入PyTorch的MNIST数据集
import numpy as np  # 导入NumPy库
import matplotlib.pyplot as plt  # 导入Matplotlib库，用于数据可视化
import matplotlib.gridspec as gridspec  # 导入Matplotlib的网格规格模块，用于定制图像的布局

# 设置matplotlib的默认图形大小为10.0x8.0
plt.rcParams['figure.figsize'] = (10.0, 8.0)
# 设置图像插值方式为最近邻插值
plt.rcParams['image.interpolation'] = 'nearest'
# 设置图像的颜色映射为灰度
plt.rcParams['image.cmap'] = 'gray'


def show_images(images):  # 定义画图工具
    """
    显示一组图像
    参数:
    images: 二维数组，包含多个图像数据
    """
    # 将图像数据重塑为二维数组，方便后续操作
    images = np.reshape(images, [images.shape[0], -1])
    # 计算图像的行数和列数，以便确定画布的大小
    sqrtn = int(np.ceil(np.sqrt(images.shape[0])))
    sqrtimg = int(np.ceil(np.sqrt(images.shape[1])))
    # 创建一个新的画布，大小为图像的行数和列数
    fig = plt.figure(figsize=(sqrtn, sqrtn))
    # 创建一个网格布局，用于在画布上放置图像
    gs = gridspec.GridSpec(sqrtn, sqrtn)
    # 更新网格布局的间距
    gs.update(wspace=0.05, hspace=0.05)
    # 遍历图像数据
    for i, img in enumerate(images):
        # 在画布上创建一个子图
        ax = plt.subplot(gs[i])
        # 关闭子图的坐标轴标签
        plt.axis('off')
        # 移除子图的x轴标签
        ax.set_xticklabels([])
        # 移除子图的y轴标签
        ax.set_yticklabels([])
        # 设置子图的纵横比一致
        ax.set_aspect('equal')
        # 在子图上显示图像
        plt.imshow(img.reshape([sqrtimg, sqrtimg]))
    # 显示所有图像
    plt.show()


def preprocess_img(x):
    """
    对输入的图像进行预处理，将图像数据转换为Tensor类型，并进行归一化处理

    参数:
    x: 输入的图像数据

    返回值:
    返回预处理后的图像数据
    """
    # 将图像数据转换为Tensor类型
    x = tfs.ToTensor()(x)
    # 对Tensor类型的图像数据进行归一化处理，使其范围在-1到1之间
    return (x - 0.5) / 0.5


def deprocess_img(x):
    """
    对输入的图像进行反预处理，将归一化的图像数据恢复到原始范围

    参数:
    x: 输入的归一化图像数据

    返回值:
    返回恢复到原始范围的图像数据
    """
    # 对归一化的图像数据进行反归一化处理，使其范围在0到1之间
    return (x + 1.0) / 2.0


# 代码3-36
# 定义一个取样的函数
class ChunkSampler(sampler.Sampler):

    def __init__(self, num_samples, start=0):
        """
        初始化ChunkSampler对象
        :param num_samples: 要采样的数量
        :param start: 采样开始的位置，默认为0
        """
        # 采样数量
        self.num_samples = num_samples
        # 采样开始位置
        self.start = start

    def __iter__(self):
        """
        返回一个迭代器，用于生成采样索引
        :return: 迭代器
        """
        return iter(range(self.start, self.start + self.num_samples))

    def __len__(self):
        """
        返回采样数量
        :return: 采样数量
        """
        return self.num_samples

# 定义训练集样本数量
NUM_TRAIN = 50000
# 定义验证集样本数量
NUM_VAL = 5000
# 定义噪声维度
NOISE_DIM = 96
# 定义批次大小
batch_size = 128
# 创建训练数据集
train_set = MNIST('data3/', train=True, download=False, transform=preprocess_img)
# 使用ChunkSampler采样器加载训练数据集
train_data = DataLoader(train_set, batch_size=batch_size, sampler=ChunkSampler(NUM_TRAIN, 0))
# 创建验证数据集
val_set = MNIST('data3/', train=True, download=False, transform=preprocess_img)
# 使用ChunkSampler采样器加载验证数据集
val_data = DataLoader(val_set, batch_size=batch_size, sampler=ChunkSampler(NUM_VAL, NUM_TRAIN))
# 可视化图片效果
imgs = deprocess_img(train_data.__iter__().next()[0].view(batch_size, 784)).numpy().squeeze()
# 显示图片
show_images(imgs)

# 代码3-37
# 定义判别器
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 定义设备，优先使用cuda，否则使用cpu
class build_dc_classifier(nn.Module):  # 定义判别器类

    def __init__(self):  # 初始化函数
        super(build_dc_classifier, self).__init__()  # 调用父类的初始化函数
        self.conv = nn.Sequential(  # 定义卷积层
            nn.Conv2d(1, 32, 5, 1),  # 输入通道数为1，输出通道数为32，卷积核大小为5，步长为1
            nn.LeakyReLU(0.01),  # 定义LeakyReLU激活函数，负斜率为0.01
            nn.MaxPool2d(2, 2),  # 定义最大池化层，池化窗口大小为2，步长为2
            nn.Conv2d(32, 64, 5, 1),  # 输入通道数为32，输出通道数为64，卷积核大小为5，步长为1
            nn.LeakyReLU(0.01),  # 定义LeakyReLU激活函数，负斜率为0.01
            nn.MaxPool2d(2, 2))  # 定义最大池化层，池化窗口大小为2，步长为2

        self.fc = nn.Sequential(  # 定义全连接层
            nn.Linear(1024, 1024),  # 输入特征数为1024（4*4*64=1024），输出特征数为1024
            nn.LeakyReLU(0.01),  # 定义LeakyReLU激活函数，负斜率为0.01
            nn.Linear(1024, 1))  # 输入特征数为1024，输出特征数为1

    def forward(self, x):  # 前向传播函数
        x = self.conv(x)  # 卷积层处理输入x
        x = x.view(x.shape[0], -1)  # 将x展平
        x = self.fc(x)  # 全连接层处理展平后的x
        return x  # 返回处理结果


# 代码3-38
# 定义卷积生成网络
class build_dc_generator(nn.Module):  # 定义生成器类

    def __init__(self, noise_dim=NOISE_DIM):  # 初始化函数
        super(build_dc_generator, self).__init__()  # 调用父类的初始化函数
        self.fc = nn.Sequential(  # 定义全连接层
            nn.Linear(noise_dim, 1024),  # 输入特征数为噪声维度，输出特征数为1024
            nn.ReLU(True),  # 定义ReLU激活函数
            nn.BatchNorm1d(1024),  # 定义批量归一化层
            nn.Linear(1024, 7 * 7 * 128),  # 输入特征数为1024，输出特征数为7*7*128
            nn.ReLU(True),  # 定义ReLU激活函数
            nn.BatchNorm1d(7 * 7 * 128))  # 定义批量归一化层

        self.conv = nn.Sequential(  # 定义卷积转置层
            nn.ConvTranspose2d(128, 64, 4, 2, padding=1),  # 输入通道数为128，输出通道数为64，卷积核大小为4，步长为2，填充为1
            nn.ReLU(True),  # 定义ReLU激活函数
            nn.BatchNorm2d(64),  # 定义批量归一化层
            nn.ConvTranspose2d(64, 1, 4, 2, padding=1))  # 输入通道数为64，输出通道数为1，卷积核大小为4，步长为2，填充为1

    def forward(self, x):  # 前向传播函数
        x = self.fc(x)  # 全连接层处理输入x
        x = x.view(x.shape[0], 128, 7, 7)  # 将x展平，通道数为128，大小为7x7
        x = self.conv(x)  # 卷积转置层处理展平后的x
        return x  # 返回处理结果


# 代码3-39
# 定义损失函数
bce_loss = nn.BCEWithLogitsLoss()  # 定义二元交叉熵损失函数
def discriminator_loss(logits_real, logits_fake):  # 判别器的loss
    size = logits_real.shape[0]  # 获取输入数据的数量
    true_labels = Variable(torch.ones(size, 1)).float().to(device)  # 真实数据的标签为1
    false_labels = Variable(torch.zeros(size, 1)).float().to(device)  # 假数据的标签为0
    loss = bce_loss(logits_real, true_labels) + bce_loss(logits_fake, false_labels)  # 计算判别器的loss
    return loss  # 返回判别器的loss

def generator_loss(logits_fake):  # 生成器的loss
    size = logits_fake.shape[0]  # 获取输入数据的数量
    true_labels = Variable(torch.ones(size, 1)).float().to(device)  # 假数据的标签为1
    loss = bce_loss(logits_fake, true_labels)  # 计算生成器的loss
    return loss  # 返回生成器的loss

# 代码3-44
# 构建网络
def train_dc_gan(D_net,  # 判别网络
              G_net,  # 生成网络
              D_optimizer,  # 判别网络优化器
              G_optimizer,  # 生成网络优化器
              discriminator_loss,  # 判别器的loss函数
              generator_loss,  # 生成器的loss函数
              show_every=250,  # 每250次迭代显示一次
              noise_size=96,  # 噪声维度
              num_epochs=10):  # 训练周期数
    iter_count = 0  # 迭代次数初始化
    for epoch in range(num_epochs):  # 遍历训练周期数
        for x, _ in train_data:  # 遍历训练数据
            bs = x.shape[0]  # 获取输入数据的数量
            real_data = Variable(x).to(device)  # 真实数据
            logits_real = D_net(real_data)  # 判别网络得分
            # -1~1的均匀分布
            sample_noise = (torch.rand(bs, noise_size) - 0.5) / 0.5
            g_fake_seed = Variable(sample_noise).to(device)
            fake_images = G_net(g_fake_seed)  # 生成的假的数据
            logits_fake = D_net(fake_images)  # 判别网络得分
            # 判别器的loss
            d_total_error = discriminator_loss(logits_real, logits_fake)
            D_optimizer.zero_grad()  # 判别网络优化器梯度清零
            d_total_error.backward()  # 反向传播计算判别网络的梯度
            D_optimizer.step()  # 更新判别网络的参数
            # 生成网络
            g_fake_seed = Variable(sample_noise).to(device)
            fake_images = G_net(g_fake_seed)  # 生成的假的数据
            gen_logits_fake = D_net(fake_images)
            g_error = generator_loss(gen_logits_fake)  # 生成网络的 loss
            G_optimizer.zero_grad()  # 生成网络优化器梯度清零
            g_error.backward()  # 反向传播计算生成网络的梯度
            G_optimizer.step()  # 更新生成网络的参数
            if (iter_count % show_every == 0):  # 每250次迭代显示一次
                print('Iter: {}, D: {:.4}, G:{:.4}'.format(iter_count,
                                                           d_total_error.item(),
                                                           g_error.item()))
                imgs_numpy = deprocess_img(fake_images.data.cpu().numpy())
                show_images(imgs_numpy[0:16])
                plt.show()
                print()
            iter_count += 1

# 代码3-45
# 定义优化器
def get_optimizer(net):
    """
    定义优化器，使用adam优化器，学习率为3e-4，beta1为0.5，beta2为0.999
    :param net: 输入的网络模型
    :return: 返回定义好的优化器
    """
    optimizer = torch.optim.Adam(net.parameters(), lr=3e-4, betas=(0.5, 0.999))
    return optimizer


# 代码3-46
# 训练网络
# 构建判别器和生成器
D_DC = build_dc_classifier().to(device)  # 构建判别器
G_DC = build_dc_generator().to(device)  # 构建生成器
# 创建判别器的优化器
D_DC_optim = get_optimizer(D_DC)  # 判别器创建优化器
# 创建生成器的优化器
G_DC_optim = get_optimizer(G_DC)  # 生成器创建优化器
# 训练判别器和生成器
train_dc_gan(D_DC, G_DC, D_DC_optim, G_DC_optim,
           discriminator_loss, generator_loss, num_epochs=20)  # 训练20个周期